{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a10bdb46",
   "metadata": {},
   "source": [
    "# Memory Collections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52531739",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "\n",
    "class InteractionDirective(BaseModel):\n",
    "    directive: str = Field(\n",
    "        description=(\n",
    "            \"An independent, atomic instruction that describes one specific way the assistant should adapt its interaction to the user. \"\n",
    "            \"Each directive should focus on a single behavior or style — do not combine multiple preferences into one. \"\n",
    "            \"Directives must be phrased as imperatives. \"\n",
    "            \"Examples: \"\n",
    "            \"'Explain technical topics using a step-by-step format.', \"\n",
    "            \"'Use concise code examples.', \"\n",
    "            \"'Avoid adding explanations unless explicitly requested.', \"\n",
    "            \"'Include inline comments in code when requested.', \"\n",
    "            \"'Ask if the user wants a deeper explanation after giving an answer.'\"\n",
    "        )\n",
    "    )\n",
    "\n",
    "class InteractionStyle(BaseModel):\n",
    "    directives: list[InteractionDirective] = Field(\n",
    "        description=\"A collection of interaction-relevant insights to guide how the assistant should respond to this user.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fa7eafa",
   "metadata": {},
   "source": [
    "## Creating collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f0d0894",
   "metadata": {},
   "outputs": [],
   "source": [
    "from trustcall import create_extractor\n",
    "from langchain_core.messages import HumanMessage, AIMessage\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "# trustcall_extractor = create_extractor(model, tools=[UserFacts])\n",
    "\n",
    "trustcall_extractor = create_extractor(\n",
    "    model,\n",
    "    tools=[InteractionStyle],\n",
    "    tool_choice=\"InteractionStyle\"\n",
    ")\n",
    "\n",
    "\n",
    "# Invoke the extractor\n",
    "result = trustcall_extractor.invoke([HumanMessage(content=\"\"\"\n",
    "Hey, how are you? I am Evgeny, I am a junior software engineer and I need your help with \n",
    "my project. This is a TaskManager Java Spring Boot application. I do not get how\n",
    "to configure security there!\n",
    "Also, I often come back to these answers later, so please include comments inside any code you share.\n",
    "\"\"\")])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73e1b7d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in result[\"messages\"]:\n",
    "    m.pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6437fa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "result[\"responses\"][0].directives"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4095b02",
   "metadata": {},
   "source": [
    "### let's save initial memories into long-term memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b44c914",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.store.memory import InMemoryStore\n",
    "\n",
    "# Initialize the in-memory store\n",
    "in_memory_store = InMemoryStore()\n",
    "\n",
    "# Namespace for the memory to save\n",
    "user_id = \"1\"\n",
    "namespace_for_directives = (user_id, \"directives\")\n",
    "\n",
    "for i, directive in enumerate(result[\"responses\"][0].directives):\n",
    "    in_memory_store.put(namespace_for_directives, f\"directive-{i}\", directive.model_dump())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf664f9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "for directive in in_memory_store.search(namespace_for_directives):\n",
    "    print(directive.dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de93c393",
   "metadata": {},
   "source": [
    "## Updating collection\n",
    "\n",
    "The same as for user profile we want to be able to update the collection, not recreate it everytime. So what we need:\n",
    "\n",
    "1. update already existing memory items if there is an update for the fact\n",
    "2. create new memory items if there is a new fact\n",
    "\n",
    "a way how to do this with Trustcall is to use `enable_inserts=True` and single elements Schema. Another thing we have to solve - to provide memory ids for update. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "878ef8b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from trustcall import create_extractor\n",
    "\n",
    "# Create the extractor\n",
    "trustcall_extractor = create_extractor(\n",
    "    model,\n",
    "    tools=[InteractionDirective],\n",
    "    tool_choice=\"InteractionDirective\",\n",
    "    enable_inserts=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68f59c63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare existing facts\n",
    "schema = \"InteractionDirective\"\n",
    "existing_directives = in_memory_store.search(namespace_for_directives)\n",
    "memories = (\n",
    "    [(existing_directive.key, schema, existing_directive.value) for existing_directive in existing_directives] \n",
    "    if existing_directives else None\n",
    ")\n",
    "memories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5319af4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage, AIMessage\n",
    "\n",
    "conversation = [\n",
    "    HumanMessage(content=\"\"\"\n",
    "    Hey, how are you? I’m Evgeny, I’m a software engineer and I need your help with my project.\n",
    "    This is a TaskManager Java Spring Boot application. I do not get how to configure security there!\n",
    "    Also, I often come back to these answers later, so please include comments inside any code you share.\n",
    "    \"\"\"), \n",
    "\n",
    "    AIMessage(content=\"\"\"\n",
    "    Hey Evgeny! I’d be glad to help. Spring Security can be a bit tricky at first. Do you want a quick step-by-step setup example, \n",
    "    or would you prefer a deeper explanation of how it works?\n",
    "    \"\"\"), \n",
    "\n",
    "    HumanMessage(content=\"\"\"\n",
    "    Sorry if I sound a bit stressed — I’m not very experienced with Spring Security.\n",
    "    If you could break things down step by step, that would really help.\n",
    "    \"\"\"),\n",
    "\n",
    "    AIMessage(content=\"\"\"\n",
    "    Totally understood! No worries at all — I’ll walk you through it one step at a time.\n",
    "    Let’s start with the basic configuration: you’ll need to add the Spring Security dependency first. Here's how...\n",
    "    \"\"\"), \n",
    "\n",
    "    HumanMessage(content=\"\"\"\n",
    "    Thanks, but could you skip the long explanation? Just tell me exactly what I need to do.\n",
    "    \"\"\"),\n",
    "]\n",
    "\n",
    "# Update the instruction\n",
    "system_msg = \"\"\"\n",
    "Update existing directives about the user and create new ones based on the following conversation, \n",
    "ensuring each directive reflects one specific communication preference or behavior the assistant should follow.\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "result = trustcall_extractor.invoke({\"messages\": [system_msg] + conversation, \n",
    "                                     \"existing\": memories})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00fb98a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Messages from the model indicate two tool calls were made\n",
    "for m in result[\"messages\"]:\n",
    "    m.pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14c9c0c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "result[\"response_metadata\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae645d7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print the results\n",
    "print(\"Updated and new person records:\")\n",
    "for r, rmeta in zip(result[\"responses\"], result[\"response_metadata\"]):\n",
    "    print(f\"ID: {rmeta.get('json_doc_id', 'New')}\")\n",
    "    print(r.model_dump_json(indent=2))\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96fb3189",
   "metadata": {},
   "source": [
    "## ChatBot with updateable directions from memory (in studio)\n",
    "\n",
    "see [studio/directive_memory_bot.py](studio/directive_memory_bot.py)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
